import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyPairGenerator;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumPrivateKeyParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumPublicKeyParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumSigner;
import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider;
import org.bouncycastle.util.io.pem.PemObject;
import org.bouncycastle.util.io.pem.PemWriter;

import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.math.BigInteger;
import java.security.Security;
import java.security.SecureRandom;
import java.security.cert.X509Certificate;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Date;

public class JcaDilithiumCertExample2 {
    static {
        Security.addProvider(new BouncyCastleProvider());
        Security.addProvider(new BouncyCastlePQCProvider());
    }

    public static void main(String[] args) throws Exception {
        try {
            System.out.println("使用底层Dilithium API生成证书...");
            generateCertificateWithLowLevelAPI();
            
        } catch (Exception e) {
            System.out.println("证书生成失败: " + e.getMessage());
            e.printStackTrace();
        }
    }

    private static void generateCertificateWithLowLevelAPI() throws Exception {
        SecureRandom random = new SecureRandom();
        
        // 使用Dilithium5参数
        DilithiumParameters params = DilithiumParameters.dilithium5;
        DilithiumKeyPairGenerator keyPairGen = new DilithiumKeyPairGenerator();
        keyPairGen.init(new org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyGenerationParameters(random, params));
        
        var keyPair = keyPairGen.generateKeyPair();
        DilithiumPublicKeyParameters pubKey = (DilithiumPublicKeyParameters) keyPair.getPublic();
        DilithiumPrivateKeyParameters privKey = (DilithiumPrivateKeyParameters) keyPair.getPrivate();

        System.out.println("Dilithium5密钥对生成成功");

        // 设置正确的OID（根据NIST标准）
        String algorithmOID = getAlgorithmOID(params);
        System.out.println("使用算法OID: " + algorithmOID);

        // 构建证书
        X500Name subject = new X500Name("CN=Test Dilithium Certificate, O=Demo, C=CN");
        BigInteger serial = BigInteger.valueOf(System.currentTimeMillis());
        Instant now = Instant.now();
        Date notBefore = Date.from(now);
        Date notAfter = Date.from(now.plus(365, ChronoUnit.DAYS));

        // 算法标识符
        AlgorithmIdentifier sigAlgId = new AlgorithmIdentifier(new ASN1ObjectIdentifier(algorithmOID));
        AlgorithmIdentifier keyAlgId = new AlgorithmIdentifier(new ASN1ObjectIdentifier(algorithmOID));

        // 公钥信息
        SubjectPublicKeyInfo pubInfo = new SubjectPublicKeyInfo(keyAlgId, pubKey.getEncoded());

        X509v3CertificateBuilder certGen = new X509v3CertificateBuilder(
            subject, serial, notBefore, notAfter, subject, pubInfo);

        // 自定义内容签名器
        ContentSigner signer = new ContentSigner() {
            private final java.io.ByteArrayOutputStream buffer = new java.io.ByteArrayOutputStream();
            
            @Override
            public AlgorithmIdentifier getAlgorithmIdentifier() {
                return sigAlgId;
            }
            
            @Override
            public java.io.OutputStream getOutputStream() {
                return buffer;
            }
            
            @Override
            public byte[] getSignature() {
                try {
                    byte[] dataToSign = buffer.toByteArray();
                    DilithiumSigner dilithiumSigner = new DilithiumSigner();
                    dilithiumSigner.init(true, privKey);
                    return dilithiumSigner.generateSignature(dataToSign);
                } catch (Exception e) {
                    throw new RuntimeException("签名失败: " + e.getMessage(), e);
                }
            }
        };

        System.out.println("开始构建证书...");
        org.bouncycastle.cert.X509CertificateHolder certHolder = certGen.build(signer);
        System.out.println("证书构建完成");

        // 转换为JCA证书
        X509Certificate cert = new JcaX509CertificateConverter()
            .setProvider("BC")
            .getCertificate(certHolder);

        System.out.println("证书生成成功");

        // 输出文件
        writeCertificateToPEM(cert, "dilithium_cert.crt");
        writePrivateKeyToPEM(privKey.getEncoded(), "dilithium_private_key.pem");
        writePublicKeyToPEM(pubKey.getEncoded(), "dilithium_public_key.pem");
        writeCertificateToDER(cert, "dilithium_cert.der");

        System.out.println("=== 证书信息 ===");
        System.out.println("主题: " + cert.getSubjectX500Principal());
        System.out.println("颁发者: " + cert.getIssuerX500Principal());
        System.out.println("签名算法: " + cert.getSigAlgName());
        System.out.println("序列号: " + cert.getSerialNumber());
        System.out.println("有效期: " + cert.getNotBefore() + " 至 " + cert.getNotAfter());
        System.out.println("所有文件已成功生成！");
    }

    private static String getAlgorithmOID(DilithiumParameters params) {
        // 根据Dilithium参数返回正确的OID
        if (params == DilithiumParameters.dilithium2) {
            return "2.16.840.1.101.3.4.3.17"; // Dilithium2 OID
        } else if (params == DilithiumParameters.dilithium3) {
            return "2.16.840.1.101.3.4.3.18"; // Dilithium3 OID
        } else if (params == DilithiumParameters.dilithium5) {
            return "2.16.840.1.101.3.4.3.19"; // Dilithium5 OID
        } else {
            return "2.16.840.1.101.3.4.3.19"; // 默认使用Dilithium5 OID
        }
    }

    // PEM输出方法
    private static void writeCertificateToPEM(X509Certificate cert, String filename) throws Exception {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("CERTIFICATE", cert.getEncoded());
            pemWriter.writeObject(pemObject);
        }
        System.out.println("证书已保存: " + filename);
    }

    private static void writePrivateKeyToPEM(byte[] privateKeyBytes, String filename) throws IOException {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("PRIVATE KEY", privateKeyBytes);
            pemWriter.writeObject(pemObject);
        }
        System.out.println("私钥已保存: " + filename);
    }

    private static void writePublicKeyToPEM(byte[] publicKeyBytes, String filename) throws IOException {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("PUBLIC KEY", publicKeyBytes);
            pemWriter.writeObject(pemObject);
        }
        System.out.println("公钥已保存: " + filename);
    }

    private static void writeCertificateToDER(X509Certificate cert, String filename) throws Exception {
        try (FileOutputStream fos = new FileOutputStream(filename)) {
            fos.write(cert.getEncoded());
        }
        System.out.println("DER证书已保存: " + filename);
    }

    // 生成不同安全级别的证书
    public static void generateDifferentLevels() throws Exception {
        DilithiumParameters[] levels = {
            DilithiumParameters.dilithium2,
            DilithiumParameters.dilithium3,
            DilithiumParameters.dilithium5
        };
        
        for (DilithiumParameters level : levels) {
            try {
                System.out.println("\n生成 " + level.getName() + " 证书...");
                generateSpecificLevel(level);
            } catch (Exception e) {
                System.out.println(level.getName() + " 证书生成失败: " + e.getMessage());
            }
        }
    }

    private static void generateSpecificLevel(DilithiumParameters params) throws Exception {
        SecureRandom random = new SecureRandom();
        DilithiumKeyPairGenerator keyPairGen = new DilithiumKeyPairGenerator();
        keyPairGen.init(new org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyGenerationParameters(random, params));
        
        var keyPair = keyPairGen.generateKeyPair();
        DilithiumPublicKeyParameters pubKey = (DilithiumPublicKeyParameters) keyPair.getPublic();
        DilithiumPrivateKeyParameters privKey = (DilithiumPrivateKeyParameters) keyPair.getPrivate();

        String algorithmOID = getAlgorithmOID(params);
        
        X500Name subject = new X500Name("CN=Test " + params.getName() + " Certificate, O=Demo, C=CN");
        BigInteger serial = BigInteger.valueOf(System.currentTimeMillis());
        Instant now = Instant.now();
        Date notBefore = Date.from(now);
        Date notAfter = Date.from(now.plus(365, ChronoUnit.DAYS));

        AlgorithmIdentifier sigAlgId = new AlgorithmIdentifier(new ASN1ObjectIdentifier(algorithmOID));
        AlgorithmIdentifier keyAlgId = new AlgorithmIdentifier(new ASN1ObjectIdentifier(algorithmOID));

        SubjectPublicKeyInfo pubInfo = new SubjectPublicKeyInfo(keyAlgId, pubKey.getEncoded());

        X509v3CertificateBuilder certGen = new X509v3CertificateBuilder(
            subject, serial, notBefore, notAfter, subject, pubInfo);

        ContentSigner signer = new ContentSigner() {
            private final java.io.ByteArrayOutputStream buffer = new java.io.ByteArrayOutputStream();
            
            @Override
            public AlgorithmIdentifier getAlgorithmIdentifier() { return sigAlgId; }
            
            @Override
            public java.io.OutputStream getOutputStream() { return buffer; }
            
            @Override
            public byte[] getSignature() {
                byte[] dataToSign = buffer.toByteArray();
                DilithiumSigner dilithiumSigner = new DilithiumSigner();
                dilithiumSigner.init(true, privKey);
                return dilithiumSigner.generateSignature(dataToSign);
            }
        };

        X509Certificate cert = new JcaX509CertificateConverter()
            .setProvider("BC")
            .getCertificate(certGen.build(signer));

        writeCertificateToPEM(cert, params.getName().toLowerCase() + "_cert.crt");
        writePrivateKeyToPEM(privKey.getEncoded(), params.getName().toLowerCase() + "_private.pem");
        
        System.out.println(params.getName() + " 证书生成成功");
    }
}